import httpx
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response
import logging

# 创建全局HTTP客户端，启用连接池
client = httpx.AsyncClient(
    base_url="http://localhost:5173",
    timeout=30.0,
    limits=httpx.Limits(max_keepalive_connections=5, max_connections=10)
)

class ViteProxyMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, callNext):
        # 如果是API请求，直接传递给FastAPI
        if request.url.path.startswith("/api"):
            return await callNext(request)

        try:
            # 构建目标URL
            targetUrl = f"http://localhost:5173{request.url.path}"
            if request.url.query:
                targetUrl += f"?{request.url.query}"

            # 转发请求
            resp = await client.request(
                method=request.method,
                url=targetUrl,
                headers=dict(request.headers),
                content=await request.body(),
                follow_redirects=True
            )

            # 构建响应
            return Response(
                content=resp.content,
                status_code=resp.status_code,
                headers=dict(resp.headers),
                media_type=resp.headers.get("content-type")
            )
        except httpx.ConnectError:
            logging.error("无法连接到Vite开发服务器，请确保Vite正在运行")
            return Response(
                content="Vite开发服务器未运行",
                status_code=503
            )
        except Exception as e:
            logging.error(f"代理请求时发生错误: {str(e)}")
            return Response(
                content=f"代理错误: {str(e)}",
                status_code=500
            )

# 导出中间件函数
def viteProxyMiddleware(app):
    return ViteProxyMiddleware(app=app) 